#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# standard
import os
import sys
import json
import threading
import ctypes
from datetime import datetime
from typing import Callable
from datetime import datetime
from typing import Callable
from typing import List
from typing import Iterable
from dataclasses import dataclass
import time

# third
import loguru
from pika.connection import URLParameters
from pika import BlockingConnection
from pika import BasicProperties
from pika.adapters.blocking_connection import BlockingChannel
from pika.spec import Basic
from pika.exceptions import *
import requests
import urllib3
from urllib.parse import quote_plus

# local
_P_PATH =  os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
if _P_PATH not in sys.path:
    sys.path.append(_P_PATH)
from models import *


######## rabbitmq 相关函数#########
######## 代码开始 ##############
# pika 文档地址 https://pika.readthedocs.io/en/stable/
# 文件的新旧 zabbix-receiver == aiops后台，== web-extend-apis == mq-message-listener



urllib3.disable_warnings()
USER_NAME = "developer"
PASSWORD = "developer123"
HOST = "61.169.129.194"
PORT = 18104
V_HOST = ""
HEART_BEAT = 600
TIMEOUT = 300
URL = f"amqp://{USER_NAME}:{quote_plus(PASSWORD)}@{HOST}:{PORT}/{V_HOST}?heartbeat={HEART_BEAT}&blocked_connection_timeout={TIMEOUT}"
PARAMS = URLParameters(URL)
PROJECT_NAME = os.path.basename(_P_PATH).replace("-", "_")
ROUTING_KEY = "zabbix7_alert"


@dataclass
class ConsumerSetting:
    """
    消费设置
    """
    exchange: str = "zabbix.topic"  # 交换机名称
    queue: str = f"{PROJECT_NAME}.{ENV}.{ROUTING_KEY}"  # 队列名称
    routing_key: str = ROUTING_KEY  # 路由键
    callback: Callable = None  # 回调函数

    def __init__(self, exchange: str = None, queue: str = None, routing_key: str = None, callback: Callable = None):
        super().__init__()
        if exchange:
            self.exchange = exchange
        if queue:
            self.queue = queue
        if routing_key:
            self.routing_key = routing_key
        if callback:
            self.callback = callback


@dataclass
class RabbitMQManager:
    """
    RabbitMQ管理器
    队列名的组成
    三段式： 项目名 + 环境 + routing_key
    {PROJECT_NAME}.{ENV}.{ROUTING_KEY}
    """
    exchange: str = "zabbix.topic"  # 交换机名称
    routing_key: str = ROUTING_KEY  # 路由键
    queue: str = f"{PROJECT_NAME}.{ENV}.{ROUTING_KEY}"  # 队列名称
    expiration: int = 60  # 消息的有效期,单位秒
    connect_fail_count: int = 0  # 连接失败计数
    connect_max_retry: int = 10  # 最大重试次数
    connect_retry_interval: int = 1  # 重试间隔,单位秒
    connection: BlockingConnection = None   # RabbitMq的阻塞式客户端
    channel: BlockingChannel = None  # RabbitMq的阻塞式通道
    callback: Callable = None  # callback函数的优先级顺序 RabbitMQManager.settings.ConsumerSetting.callback > RabbitMQManager.callback > RabbitMQManager.default_callback
    settings: List[ConsumerSetting] = None  # 消费设置

    def __init__(self, settings: List[ConsumerSetting] = None, callback: Callable = None,
                 routing_key: str = ROUTING_KEY, exchange: str = None, 
                 expiration: int = None):
        super().__init__()
        if settings:
            self.settings = settings
        if callback:
            self.callback = callback
        if routing_key:
            self.routing_key = routing_key
            self.queue = f"{PROJECT_NAME}.{routing_key}.{ENV}"
        if exchange:
            self.exchange = exchange
        if expiration:
            self.expiration = expiration
        self.connection = self.check_mq_connect()

    def get_connection(self) -> BlockingConnection:
        """
        返回一个RabbitMq的阻塞式客户端
        :return:
        """
        try:
            self.connection = BlockingConnection(PARAMS)
        except AMQPConnectionError as e:
            logger.error(e)
            self.connection = None
        except Exception as e:
            raise e
        return self.connection

    @staticmethod
    def send_message_base(channel: BlockingChannel, payload: dict, exchange: str, routing_key: str = None, expiration: int | None = None):
        """
        发送消息,这是一个比较低级的方法,调用这需要自己提供一个Channel,
        你需要
        1. 自行声明交换机和队列
        2. 自行关闭Channel和Connection,
        否则可能会导致服务器资源耗尽
        如果你只是临时发送消息.建议使用 send_once函数.
        :param channel: 通道
        :param payload: 发送消息的负载,字典格式,要求必须能进行json序列化.
        :param exchange: 交换机名称. 空字符表示默认队列
        :param routing_key:
        :param expiration: 消息的有效期,单位秒.
        """
        if isinstance(expiration, int):
            expiration = str(expiration * 1000)
        else:
            expiration = None
        properties = BasicProperties(
            delivery_mode=2,  # 持久化标志
            content_encoding="utf-8",
            expiration=expiration
        )
        body = json.dumps(payload)
        kw = {
            "exchange": exchange,
            "routing_key": routing_key,
            "body": body,
            "properties": properties,
        }
        channel.basic_publish(**kw)


    def send_and_close(self, routing_key: str = None, payload: dict = None, exchange: str = None,
                queue: str = None, expiration: int = None, **kwargs) -> None:
        """
        一次性发送消息,发送完毕后就会关闭连接, 不会声明交换机. 一般只在演示或者性能要求不高的场合使用。
        :param connection: RabbitMq的阻塞式客户端
        :param payload: 数据负载
        :param exchange: 交换机名称,可选的是  workflow.direct ， workflow.fanout 和 workflow.topic
        :param queue:　队列，如果你发送消息时队列还不存在，你需要填充此参数．
        :param routing_key:　路由键，必须指定．
        :param expiration:　消息默认过期时间，单位秒, 默认None, 意思是不过期
        :return:
        """
        connection = self.check_mq_connect()
        channel = connection.channel()
        routing_key = self.routing_key if routing_key is None else routing_key
        exchange = self.exchange if exchange is None else exchange  
        queue = self.queue if queue is None else queue
        expiration = self.expiration if expiration is None else expiration
        if queue is not None:
            kw1 = {
                "queue": queue,
                "durable": True,
                "auto_delete": False,
                "arguments": {
                    "x-expires": 7 * 86400 * 1000,  # 队列闲置过期时间,
                    "x-message-ttl": 1 * 86400 * 1000,  # # 设置消息的过期时间,默认是1天
                }
            }
            channel.queue_declare(**kw1)  # 声明队列
            # 以不同的参数进行相同名称的队列声明会导致channel关闭并抛出错误.
            channel.queue_bind(queue=queue, exchange=exchange)

        payload = {
            "message": "hello",
            "time": datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")
        } if not isinstance(payload, dict) else payload
        kw = {
            "channel": channel,
            "payload": payload,
            "exchange": exchange,
            "routing_key": routing_key,
            "expiration": expiration,
        }
        self.send_message_base(**kw)
        channel.close()
        connection.close()


    def check_mq_connect(self) -> BlockingConnection:
        """
        检查并返回一个阻塞式客户端
        :return: _description_
        """
        if self.connection is None or self.connection.is_closed:
            self.get_connection()
            if self.connection is None:
                self.connect_fail_count += 1
                if self.connect_fail_count >= self.connect_max_retry:
                    raise Exception(f"check_mq_connect fail 失败次数过多 {self.connect_fail_count}")
                else:
                    logger.warning(f"check_mq_connect fail, 等待{self.connect_retry_interval}秒后重新连接 {self.connect_fail_count}")
                    time.sleep(self.connect_retry_interval)
                    self.check_mq_connect()  # 递归重试
            else:  # 无需重置 self.connect_fail_count 后面会重置
                self.channel = self.connection.channel()
        if self.channel is None or self.channel.is_closed:
            self.channel = self.connection.channel()       
        return self.connection

    def _send_mq_message(self, payload: dict) -> None:
        """
        发送消息, 发送完毕后就会不关闭连接, 
        :param connection: RabbitMq的阻塞式客户端
        :param payload: 数据负载
        :param exchange: 交换机名称,可选的是  workflow.direct ， workflow.fanout 和 workflow.topic
        :param queue:　队列，如果你发送消息时队列还不存在，你需要填充此参数．
        :param routing_key:　路由键，必须指定．
        :param expiration:　消息默认过期时间，单位秒, 默认None, 意思是不过期
        :return:
        """
        self.check_mq_connect()
        # 尝试声明队列，如果同样参数的对了已存在不会抛错
        kw1 = {
            "queue": self.queue,
            "durable": True,
            "auto_delete": False,
            "arguments": {
                "x-expires": 7 * 86400 * 1000,  # 队列闲置过期时间,
                "x-message-ttl": 1 * 86400 * 1000,  # # 设置消息的过期时间,默认是1天
            }
        }
        self.channel.queue_declare(**kw1)  # 声明队列
        # 以不同的参数进行相同名称的队列声明会导致channel关闭并抛出错误.
        self.channel.queue_bind(queue=self.queue, exchange=self.exchange)
        # 准备参数和负载
        payload = {
            "message": "hello",
            "time": datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")
        } if not isinstance(payload, dict) else payload
        kw = {
            "channel": self.channel,
            "payload": payload,
            "exchange": self.exchange,
            "routing_key": self.routing_key,
            "expiration": self.expiration,
        }
        self.send_message_base(**kw)
        logger.debug(f"发送MQ消息成功: {payload}")
        if self.connect_fail_count != 0:
            self.connect_fail_count = 0

    def send_mq_message(self, payload: dict) -> None:
        """
        发送消息,发送完毕后就会不关闭连接， 如果连接失效，会自动重建连接
        :param connection: RabbitMq的阻塞式客户端
        :param payload: 数据负载
        :param exchange: 交换机名称,可选的是  workflow.direct ， workflow.fanout 和 workflow.topic
        :param queue:　队列，如果你发送消息时队列还不存在，你需要填充此参数．
        :param routing_key:　路由键，必须指定．
        :param expiration:　消息默认过期时间，单位秒, 默认None, 意思是不过期
        :return:
        """
        try:
            self._send_mq_message(payload)
        except (AMQPConnectionError, StreamLostError, ConnectionResetError, 
                ConnectionRefusedError, ConnectionOpenAborted,
                ConnectionClosed, ConnectionBlockedTimeout) as e:       
            self.connect_fail_count += 1
            logger.error(f"发送消息失败: {e}")
            if not self.channel.is_closed:
                self.channel.close()
            if not self.connection.is_closed:
                self.connection.close()
            if self.connect_fail_count < self.connect_max_retry:
                logger.warning(f"{self.connect_retry_interval}秒后尝试重连: {self.connect_fail_count}")
                time.sleep(self.connect_retry_interval)
                self._send_mq_message(payload)
            else:
                logger.error(f"重连次数超过最大次数: {self.connect_max_retry}")
        except Exception as e:
            raise e

    @staticmethod
    def default_cb(ch: BlockingChannel, deliver: Basic.Deliver, properties: BasicProperties, body: bytes):
        """
        默认回调函数.可以作为回调函数的例子.
        :param ch: 一个rabbitmq的Channel对象.在本函数中的用途是把"我已收到消息"的消息回执传递回服务器.
        :param deliver: 传送信使.会携带传送的状态信息,用于向Channel提供消息回执
        :param properties: 消息的属性.可以查看当前文档消息来自哪个交换机，队列和routing_key等信息
        :param body: 消息体,字节格式,如果传递的时候是json字符串的传,可以直接用 json.loads函数转成字典.
        :return:
        """
        message = None
        try:
            message = json.loads(body)
        except Exception as e:
            print(e)
            message = body.decode()
        finally:
            exchange = deliver.exchange
            routing_key = deliver.routing_key
            message = "收到MQ消息: exchange = {}, routing_key = {}, message = {}".format(exchange, routing_key, message, properties.headers)
            logger.debug(message)
            ...  # your codes
            ch.basic_ack(delivery_tag=deliver.delivery_tag)  # 确认收到消息,这一步必须,否则服务端会堆积消息直至宕机.必须卸载最后一行.

    def _consume(self, settings: List[ConsumerSetting | Iterable] | None = None, callback: Callable | None = None):
        """
        消费消息, 注意,这个函数会阻塞.相当与while的状态.
        请不要在之后跟随任何业务代码.他们在本函数终止之前,不会被运行
        :param settings: 四元组组成的数组，每四元组的四个元素依次是：交换机名称，队列名称，路由键和回调函数
        :return:
        """
        self.check_mq_connect()
        if settings is None:
            if self.settings is None or len(self.settings) == 0:
                settings = [ConsumerSetting(exchange=self.exchange,
                                            queue=self.queue,
                                            routing_key=self.routing_key,
                                            callback=callback)]
            else:
                settings = self.settings
        
        auto_ack = False  # 回调函数执行失败时,会把消息返回队列
        for item in settings:
            # 设置消费信息
            if isinstance(item, ConsumerSetting):     
                exchange = item.exchange
                queue = item.queue
                routing_key = item.routing_key  
                current_callback = item.callback
            else:
                exchange = item[0]
                queue = item[1]
                routing_key = item[2]
                current_callback = item[3]
            finally_callback = current_callback if current_callback else (callback if callback else (self.callback if self.callback else self.default_cb))
            kw = {
                "queue": queue,
                "durable": True,
                "auto_delete": False,
                "arguments": {
                    "x-expires": 7 * 86400 * 1000,  # 队列闲置过期时间,
                    "x-message-ttl": 1 * 86400 * 1000,  # # 设置消息的过期时间,默认是1天
                }
            }
            self.channel.queue_declare(**kw)  # 声明队列
            self.channel.queue_bind(queue=queue, exchange=exchange, routing_key=routing_key)  # 绑定交换机
            logger.info(f"开始消费, 配置: {exchange} 队列: {queue} 路由键: {routing_key}, 回调函数: {finally_callback.__name__}")
            self.channel.basic_consume(queue=queue, on_message_callback=finally_callback, auto_ack=auto_ack)  # 消费多个队列可以重复多次本步骤      
        try:
            self.connect_fail_count = 0  # 重置连接失败次数
            self.channel.start_consuming()  # 开始消费
        except Exception as e:
            raise e
            

    def consume(self, settings: List[ConsumerSetting | Iterable] = None, callback: Callable | None = None):
        """
        消费消息, 会自动重联。注意,这个函数会阻塞.相当与while的状态.
        请不要在之后跟随任何业务代码.他们在本函数终止之前,不会被运行
        :param settings: 四元组组成的数组，每四元组的四个元素依次是：交换机名称，队列名称，路由键和回调函数
        :return:
        """
        try:
            self._consume(settings=settings, callback=callback)
        except (AMQPConnectionError, StreamLostError, ConnectionResetError, 
                ConnectionRefusedError, ConnectionOpenAborted,
                ConnectionClosed, ConnectionBlockedTimeout) as e:       
            self.connect_fail_count += 1
            logger.error(f"消费消息失败: {e}")
            if not self.channel.is_closed:
                self.channel.close()
            if not self.connection.is_closed:
                self.connection.close()
            if self.connect_fail_count < self.connect_max_retry:
                logger.warning(f"{self.connect_retry_interval}秒后尝试重连: {self.connect_fail_count}")
                time.sleep(self.connect_retry_interval)
                self.consume(settings)
            else:
                logger.error(f"重连次数超过最大次数: {self.connect_max_retry}")
        except Exception as e:
            raise e

    def disconnect(self):
        """
        断开连接
        """
        self.channel.close()
        self.connection.close()

    @classmethod
    def subscribe_mqtt_demo(cls):
        """
        订阅mqtt消息
        """
        # 接收mqtt有2个关键设置，一个是 exchange，一个是 routing_key
        # exchange 是交换机名称，这里是rabbitmq-mqtt 插件的默认值，amq.topic,
        # routing_key 是路由键，test 用来测试，也可以是其他值
        settings = [ConsumerSetting(exchange="amq.topic", routing_key="test", queue="zabbix7_alert", callback=None)]
        handler = RabbitMQManager()
        handler.consume(settings=settings)


class ConsumeThread(threading.Thread):
    """
    消费线程
    """
    settings: List[ConsumerSetting] = None
    callable: Callable | None = None

    def __init__(self, settings: List[ConsumerSetting] = None, callback: Callable | None = None):
        super().__init__(name=self.get_consume_thread_name(), daemon=True)
        self.settings = settings
        self.callable = callback

    def run(self):
        handler = RabbitMQManager(settings=self.settings, callback=self.callable)
        handler.consume()

    @staticmethod
    def get_consume_thread_name() -> str:
        """
        获取消费线程名称
        :return: _description_
        """
        
        consume_thread_name = f"mq_consumer_{PROJECT_NAME}"
        return consume_thread_name
    
    @classmethod
    def kill_consume_thread(cls):
        """
        杀死消费线程， 
        """
        consume_thread_name = cls.get_consume_thread_name()
        for t in threading.enumerate():
            if t.name == consume_thread_name:
                result = ctypes.pythonapi.PyThreadState_SetAsyncExc(t._ident, ctypes.py_object(SystemExit))  # 第一个参数时线程id，第二个参数时一个异常类的 py_object 对象
                # result 是整数， 是终结线程的数量。
                del t  # 手动清理对象更佳，之后t就变成未定义了。
                logger.warning(f"杀死消费线程 {consume_thread_name} 成功 {result}")
                break


MQSender = RabbitMQManager()  # mq 发送者


if __name__ == "__main__":
    RabbitMQManager.send_mqtt_demo({"hello": "world"})
    for i in range(100):
        MQSender.send_mq_message({"he": "343"})
        time.sleep(10)


